import pandas as pd
from datetime import datetime
from collections import defaultdict

class TradingStrategy:
    def __init__(self):
        self.max_position = 0.1  # 单行业最大仓位10%
        self.blacklist = self._load_blacklist()  # 从数据库加载黑名单
        
    def generate_daily_strategy(self, analysis_results):
        """生成每日交易策略"""
        strategy = {
            'date': datetime.now().strftime('%Y-%m-%d'),
            'long': [],
            'short': [],
            'neutral': []
        }
        
        for result in analysis_results:
            if result['sentiment']['label'] in ['POSITIVE', 4, 5]:
                strategy['long'].extend(result['related_companies'])
            elif result['sentiment']['label'] in ['NEGATIVE', 1, 2]:
                strategy['short'].extend(result['related_companies'])
            else:
                strategy['neutral'].extend(result['related_companies'])
                
        return self._apply_risk_control(strategy)
    
    def _apply_risk_control(self, raw_strategy):
        """风险控制规则"""
        # 行业分散逻辑
        industry_exposure = defaultdict(float)
        for stock in raw_strategy['long'] + raw_strategy['short']:
            industry_exposure[stock['industry']] += 1.0/len(raw_strategy['long'])
            
        # 应用仓位限制
        adjusted = []
        for stock in raw_strategy['long']:
            if industry_exposure[stock['industry']] < self.max_position:
                adjusted.append(stock)
        raw_strategy['long'] = adjusted
        
        # 过滤黑名单股票
        raw_strategy['long'] = [s for s in raw_strategy['long'] if s['code'] not in self.blacklist]
        return raw_strategy 